#include <iostream>
#include <thread>
#include <vector>
#include <atomic>
#include <chrono>
#include <cstring>
#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <sys/epoll.h>
#include <netinet/tcp.h> // For TCP_NODELAY
#include <getopt.h>

#define SERVER_IP "127.0.0.1"
#define PORT 8080

// 默认参数
int THREADS = 4;
int TOTAL_CONNECTIONS = 5000; // Reduced for easier local testing, adjust as needed
int REQUESTS_PER_CONN = 100;  // Reduced for easier local testing

std::atomic<int> total_successful_requests(0); // Renamed for clarity
std::atomic<int> total_failed_requests(0);

// Define states for a connection
enum class ConnState {
    CONNECTING,
    SENDING,
    RECEIVING
};

struct Connection {
    int sockfd;
    ConnState state;
    int requests_sent;      // Successfully sent and received a response for
    int requests_to_send;   // Total requests for this connection
    std::string send_buffer; // Buffer for the current request being sent
    size_t sent_bytes_current_request; // Bytes sent for the current request in send_buffer
    // char recv_buffer[4096]; // If you need to process response, add a receive buffer
};

void set_nonblocking(int sockfd) {
    int flags = fcntl(sockfd, F_GETFL, 0);
    fcntl(sockfd, F_SETFL, flags | O_NONBLOCK);
}

void set_tcp_nodelay(int sockfd) {
    int optval = 1;
    if (setsockopt(sockfd, IPPROTO_TCP, TCP_NODELAY, &optval, sizeof(optval)) == -1) {
        perror("setsockopt TCP_NODELAY");
        // Continue even if it fails, not critical for basic functionality
    }
}

void close_connection(int epoll_fd, Connection* conn, bool mark_remaining_failed = false) {
    if (!conn) return;
    if (conn->sockfd != -1) {
        epoll_ctl(epoll_fd, EPOLL_CTL_DEL, conn->sockfd, nullptr);
        close(conn->sockfd);
        conn->sockfd = -1; // Mark as closed
    }
    if (mark_remaining_failed) {
        int pending = conn->requests_to_send - conn->requests_sent;
        if (pending > 0) {
            total_failed_requests.fetch_add(pending, std::memory_order_relaxed);
        }
    }
    delete conn;
}

void worker(int conn_count, int req_per_conn) {
    int epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        perror("epoll_create1");
        total_failed_requests += conn_count * req_per_conn;
        return;
    }
    // std::cout << "Thread " << std::this_thread::get_id() << ": epoll_create1 success, epoll_fd: " << epoll_fd << std::endl;

    const std::string request_str = "GET / HTTP/1.1\r\nHost: localhost\r\nConnection: keep-alive\r\n\r\n";

    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    if (inet_pton(AF_INET, SERVER_IP, &server_addr.sin_addr) <= 0) {
        perror("inet_pton");
        total_failed_requests += conn_count * req_per_conn;
        close(epoll_fd);
        return;
    }

    int active_connections = 0;

    // 创建所有连接
    for (int i = 0; i < conn_count; ++i) {
        int sockfd = socket(AF_INET, SOCK_STREAM, 0); // Create blocking first for easier setsockopt
        if (sockfd == -1) {
            perror("socket");
            total_failed_requests += req_per_conn;
            continue;
        }

        set_tcp_nodelay(sockfd); // Set TCP_NODELAY before connect
        set_nonblocking(sockfd); // Then set non-blocking

        int ret = connect(sockfd, (struct sockaddr*)&server_addr, sizeof(server_addr));

        Connection* conn = new Connection();
        conn->sockfd = sockfd;
        conn->requests_sent = 0;
        conn->requests_to_send = req_per_conn;
        conn->sent_bytes_current_request = 0;
        conn->send_buffer.clear(); // Will be set when SENDING

        struct epoll_event ev;
        ev.data.ptr = conn;

        if (ret == 0) { // Connected immediately (rare for non-blocking)
            conn->state = ConnState::SENDING;
            ev.events = EPOLLOUT | EPOLLERR | EPOLLHUP | EPOLLET;
        }
        else if (errno == EINPROGRESS) {
            conn->state = ConnState::CONNECTING;
            ev.events = EPOLLOUT | EPOLLERR | EPOLLHUP | EPOLLET; // Wait for connect to complete
        }
        else {
            perror("connect");
            close(sockfd);
            delete conn;
            total_failed_requests += req_per_conn;
            continue;
        }

        if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd, &ev) == -1) {
            perror("epoll_ctl add");
            close(sockfd);
            delete conn;
            total_failed_requests += req_per_conn;
            continue;
        }
        active_connections++;
    }

    if (active_connections == 0 && conn_count > 0) {
        // All initial connection attempts failed for this thread
        close(epoll_fd);
        return;
    }


    const int MAX_EVENTS = 128; // Can be smaller per thread
    struct epoll_event events[MAX_EVENTS];

    // 事件循环
    while (active_connections > 0) {
        int n_events = epoll_wait(epoll_fd, events, MAX_EVENTS, 200); // Timeout to allow checking active_connections
        if (n_events == -1) {
            if (errno == EINTR) continue;
            perror("epoll_wait");
            break; // Exit loop on other epoll_wait errors
        }

        for (int i = 0; i < n_events; ++i) {
            struct epoll_event& ev = events[i];
            Connection* conn = static_cast<Connection*>(ev.data.ptr);
            int fd = conn->sockfd;

            if (ev.events & (EPOLLERR | EPOLLHUP)) {
                // std::cerr << "Thread " << std::this_thread::get_id() << ": EPOLLERR or EPOLLHUP on fd " << fd << std::endl;
                int so_error = 0;
                socklen_t len = sizeof(so_error);
                getsockopt(fd, SOL_SOCKET, SO_ERROR, &so_error, &len);
                // std::cerr << "SO_ERROR: " << so_error << " (" << strerror(so_error) << ")" << std::endl;

                close_connection(epoll_fd, conn, true);
                active_connections--;
                continue;
            }

            // Handle EPOLLOUT: Connection established or writable
            if (ev.events & EPOLLOUT) {
                if (conn->state == ConnState::CONNECTING) {
                    int so_error = 0;
                    socklen_t len = sizeof(so_error);
                    if (getsockopt(fd, SOL_SOCKET, SO_ERROR, &so_error, &len) < 0 || so_error != 0) {
                        // std::cerr << "Thread " << std::this_thread::get_id() << ": Connect failed on fd " << fd << " - SO_ERROR: " << so_error << " (" << strerror(so_error) << ")" << std::endl;
                        close_connection(epoll_fd, conn, true);
                        active_connections--;
                        continue;
                    }
                    // std::cout << "Thread " << std::this_thread::get_id() << ": Connection established on fd " << fd << std::endl;
                    conn->state = ConnState::SENDING;
                    // Fall through to send, or re-register if we want to be strict (usually fine to fall-through)
                }

                if (conn->state == ConnState::SENDING) {
                    if (conn->send_buffer.empty()) { // Prepare new request
                        if (conn->requests_sent >= conn->requests_to_send) {
                            // Should not happen if logic is correct, implies we finished but EPOLLOUT was still set
                           // std::cout << "Thread " << std::this_thread::get_id() << ": All requests sent for fd " << fd << ", closing." << std::endl;
                            close_connection(epoll_fd, conn); // All requests done
                            active_connections--;
                            continue;
                        }
                        conn->send_buffer = request_str;
                        conn->sent_bytes_current_request = 0;
                    }

                    while (conn->sent_bytes_current_request < conn->send_buffer.length()) {
                        ssize_t sent = send(fd,
                            conn->send_buffer.data() + conn->sent_bytes_current_request,
                            conn->send_buffer.length() - conn->sent_bytes_current_request,
                            MSG_NOSIGNAL); // MSG_NOSIGNAL to prevent SIGPIPE

                        if (sent > 0) {
                            conn->sent_bytes_current_request += sent;
                        }
                        else if (sent == 0) { // Should not happen with TCP
                            // std::cerr << "Thread " << std::this_thread::get_id() << ": send returned 0 on fd " << fd << std::endl;
                            // Treat as error or retry, for now error
                            total_failed_requests.fetch_add(1, std::memory_order_relaxed);
                            close_connection(epoll_fd, conn, true); // Mark remaining as failed
                            active_connections--;
                            goto next_event; // Break outer loop for this connection
                        }
                        else { // sent < 0
                            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                                // Send buffer is full, wait for next EPOLLOUT
                                // EPOLLET means we must keep EPOLLOUT registered
                                struct epoll_event mod_ev;
                                mod_ev.events = EPOLLOUT | EPOLLERR | EPOLLHUP | EPOLLET;
                                mod_ev.data.ptr = conn;
                                epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &mod_ev); // Ensure EPOLLOUT is still set
                                break; // Break from while loop, continue to next event in epoll_wait
                            }
                            else {
                                // perror("send error");
                                // std::cerr << "Thread " << std::this_thread::get_id() << ": Send error on fd " << fd << ": " << strerror(errno) << std::endl;
                                total_failed_requests.fetch_add(1, std::memory_order_relaxed); // Current request failed
                                close_connection(epoll_fd, conn, true); // Mark remaining as failed
                                active_connections--;
                                goto next_event; // Break outer loop for this connection
                            }
                        }
                    } // end while send loop

                    if (conn->sent_bytes_current_request == conn->send_buffer.length()) {
                        // Full request sent, switch to reading response
                        conn->send_buffer.clear();
                        conn->sent_bytes_current_request = 0;
                        conn->state = ConnState::RECEIVING;
                        struct epoll_event mod_ev;
                        mod_ev.events = EPOLLIN | EPOLLERR | EPOLLHUP | EPOLLET;
                        mod_ev.data.ptr = conn;
                        if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &mod_ev) == -1) {
                            perror("epoll_ctl mod to EPOLLIN");
                            close_connection(epoll_fd, conn, true);
                            active_connections--;
                        }
                    }
                }
            } // End EPOLLOUT

            // Handle EPOLLIN: Data available to read
            if (ev.events & EPOLLIN) {
                if (conn->state != ConnState::RECEIVING) {
                    // This shouldn't happen if state transitions are correct
                    // std::cerr << "Thread " << std::this_thread::get_id() << ": EPOLLIN event for fd " << fd << " but not in RECEIVING state." << std::endl;
                    // Potentially an error or a HUP signal disguised, close it
                    close_connection(epoll_fd, conn, true);
                    active_connections--;
                    continue;
                }

                char buffer[4096]; // Read buffer
                bool connection_ok = true;
                // For EPOLLET, we must read until EAGAIN or error
                while (true) {
                    ssize_t bytes_received = recv(fd, buffer, sizeof(buffer) - 1, 0);
                    if (bytes_received > 0) {
                        // buffer[bytes_received] = '\0'; // Null-terminate if treating as string
                        // std::cout << "Thread " << std::this_thread::get_id() << ": Received " << bytes_received << " bytes from fd " << fd << std::endl;
                        // For this benchmark, we don't process the response, just acknowledge it.
                        // If parsing HTTP, you'd accumulate data here until a full response is received.
                        // For simplicity, we assume any data means the response for one request is (starting to be) received.
                        // A more robust client would parse headers, check Content-Length, etc.
                    }
                    else if (bytes_received == 0) { // Peer closed connection
                        // std::cerr << "Thread " << std::this_thread::get_id() << ": Peer closed fd " << fd << " during recv." << std::endl;
                        total_failed_requests.fetch_add(1, std::memory_order_relaxed); // Current request failed
                        close_connection(epoll_fd, conn, true); // Mark remaining as failed
                        active_connections--;
                        connection_ok = false;
                        break;
                    }
                    else { // bytes_received < 0
                        if (errno == EAGAIN || errno == EWOULDBLOCK) {
                            // No more data to read for now
                            break; // Exit recv loop
                        }
                        else {
                            // perror("recv error");
                            // std::cerr << "Thread " << std::this_thread::get_id() << ": Recv error on fd " << fd << ": " << strerror(errno) << std::endl;
                            total_failed_requests.fetch_add(1, std::memory_order_relaxed); // Current request failed
                            close_connection(epoll_fd, conn, true); // Mark remaining as failed
                            active_connections--;
                            connection_ok = false;
                            break;
                        }
                    }
                } // End while recv loop

                if (!connection_ok) {
                    continue; // Next event from epoll_wait
                }

                // If we are here, recv loop exited due to EAGAIN or successfully read some data
                // We consider one request-response cycle complete
                total_successful_requests.fetch_add(1, std::memory_order_relaxed);
                conn->requests_sent++;

                if (conn->requests_sent < conn->requests_to_send) {
                    // More requests to send on this connection
                    conn->state = ConnState::SENDING;
                    struct epoll_event mod_ev;
                    mod_ev.events = EPOLLOUT | EPOLLERR | EPOLLHUP | EPOLLET; // Prepare to send next request
                    mod_ev.data.ptr = conn;
                    if (epoll_ctl(epoll_fd, EPOLL_CTL_MOD, fd, &mod_ev) == -1) {
                        perror("epoll_ctl mod to EPOLLOUT");
                        close_connection(epoll_fd, conn, true);
                        active_connections--;
                    }
                }
                else {
                    // All requests for this connection are done
                    // std::cout << "Thread " << std::this_thread::get_id() << ": All " << conn->requests_to_send << " requests completed for fd " << fd << ". Closing." << std::endl;
                    close_connection(epoll_fd, conn);
                    active_connections--;
                }
            } // End EPOLLIN
        next_event:;
        } // End for loop (events)
    } // End while (active_connections > 0)

    close(epoll_fd);
    // std::cout << "Thread " << std::this_thread::get_id() << " finished." << std::endl;
}

int main(int argc, char* argv[]) {
    int opt;
    while ((opt = getopt(argc, argv, "t:c:r:h")) != -1) {
        switch (opt) {
        case 't': THREADS = atoi(optarg); break;
        case 'c': TOTAL_CONNECTIONS = atoi(optarg); break;
        case 'r': REQUESTS_PER_CONN = atoi(optarg); break;
        case 'h':
        default:
            std::cerr << "Usage: " << argv[0] << " [-t threads] [-c total_connections] [-r requests_per_connection]\n";
            exit(EXIT_FAILURE);
        }
    }

    if (THREADS <= 0 || TOTAL_CONNECTIONS <= 0 || REQUESTS_PER_CONN <= 0) {
        std::cerr << "Threads, connections, and requests_per_connection must be positive integers.\n";
        exit(EXIT_FAILURE);
    }

    std::cout << "... Starting stress test with " << THREADS << " threads, "
        << TOTAL_CONNECTIONS << " total connections, and "
        << REQUESTS_PER_CONN << " requests per connection.\n"
        << "... Target: " << SERVER_IP << ":" << PORT << std::endl;

    std::vector<std::thread> threads;
    auto start_time = std::chrono::high_resolution_clock::now();

    int connections_per_thread_base = TOTAL_CONNECTIONS / THREADS;
    int remainder_connections = TOTAL_CONNECTIONS % THREADS;

    for (int i = 0; i < THREADS; ++i) {
        int conn_this_thread = connections_per_thread_base + (i < remainder_connections ? 1 : 0);
        if (conn_this_thread > 0) { // Only start thread if it has connections to manage
            threads.emplace_back(worker, conn_this_thread, REQUESTS_PER_CONN);
        }
    }

    if (threads.empty() && TOTAL_CONNECTIONS > 0) {
        std::cerr << "Warning: No worker threads started. Check if THREADS is less than TOTAL_CONNECTIONS or if calculations are correct." << std::endl;
    }


    std::cout << "... Waiting for all " << threads.size() << " worker threads to complete.\n";
    for (auto& t : threads) {
        if (t.joinable()) {
            t.join();
        }
    }

    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = end_time - start_time;
    std::cout << "... Stress test completed.\n";

    int successful = total_successful_requests.load(std::memory_order_relaxed);
    int failed = total_failed_requests.load(std::memory_order_relaxed);
    double qps = (elapsed.count() > 0) ? (successful / elapsed.count()) : 0;
    long long total_attempted = (long long)TOTAL_CONNECTIONS * REQUESTS_PER_CONN;


    std::cout << "---------------- Summary ----------------\n";
    std::cout << "Total connections attempted: " << TOTAL_CONNECTIONS << "\n";
    std::cout << "Requests per connection:   " << REQUESTS_PER_CONN << "\n";
    std::cout << "Total requests planned:    " << total_attempted << "\n";
    std::cout << "Successful requests:       " << successful << "\n";
    std::cout << "Failed requests:           " << failed << "\n";
    std::cout << "Elapsed time:              " << elapsed.count() << " seconds\n";
    std::cout << "QPS (Successful/Time):     " << qps << "\n";
    std::cout << "-----------------------------------------\n";

    if (successful + failed != total_attempted) {
        std::cout << "Warning: successful (" << successful << ") + failed (" << failed
            << ") != total_attempted (" << total_attempted << "). There might be a counting discrepancy.\n";
    }


    return 0;
}